import asyncio
import json
import ssl
from urllib.parse import urlencode
from datetime import datetime

# 引入 mcp 客户端库
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

# --- 全局禁用SSL证书验证 ---
try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    pass
else:
    ssl._create_default_https_context = _create_unverified_https_context
# --- ---


async def find_available_tickets_for_train(api_key, train_codes, departure_cities, arrival_cities, query_date=None):
    """
    使用 mcp-server-chinarailway 服务并行查询指定车次在特定日期的余票情况。
    (此版本通过 Semaphore 优化了并发控制)
    """
    # --- 【新增】定义并发控制器 ---
    # 同一时间最多允许20个查询任务在运行。这个值可以根据网络情况和服务器承受能力调整。
    CONCURRENCY_LIMIT = 20
    semaphore = asyncio.Semaphore(CONCURRENCY_LIMIT)
    # --- ---
    
    seat_map = {
        '商务座': 'swz_num', '一等座': 'zy_num', '二等座': 'ze_num', '高级软卧': 'gr_num',
        '软卧': 'rw_num', '动卧': 'dw_num', '硬卧': 'yw_num', '软座': 'rz_num',
        '硬座': 'yz_num', '无座': 'wz_num'
    }

    if isinstance(train_codes, (list, tuple, set)):
        wanted_trains = {str(t).strip() for t in train_codes if t}
    else:
        wanted_trains = {str(train_codes).strip()} if train_codes else set()

    if not query_date:
        query_date = datetime.now().strftime('%Y-%m-%d')
    
    print(f"任务开始：并行查询车次 【{', '.join(sorted(wanted_trains))}】 在日期 【{query_date}】 的余票情况。\n")
    print(f"并发限制设置为: {CONCURRENCY_LIMIT}\n")

    base_url = "https://server.smithery.ai/@other-blowsnow/mcp-server-chinarailway/mcp"
    params = {"api_key": api_key}
    url_with_auth = f"{base_url}?{urlencode(params)}"
    
    print(f"将要连接的URL (已包含认证): {base_url}")

    async def query_single_pair(session, dep_city, arr_city):
        """在同一个 session 中查询单个城市对，并受 semaphore 控制。"""
        # --- 【修改】在任务开始前获取令牌 ---
        async with semaphore:
            print(f"  [任务启动] 查询: {dep_city} -> {arr_city}")
            try:
                result = await session.call_tool(
                    "search",
                    {"date": query_date, "fromCity": dep_city, "toCity": arr_city}
                )
                
                ticket_data = []
                if isinstance(result.content, list) and len(result.content) > 0:
                    ticket_data = json.loads(result.content[0].text)
                
                for train_info in ticket_data:
                    actual_train_no = train_info.get("车次")
                    if actual_train_no and actual_train_no.strip() in wanted_trains:
                        ticket_list = train_info.get("票列表", {})
                        
                        invalid_vals = {"无", "--", "0", "", "-"}
                        available_seats = {
                            seat_name: str(raw).strip()
                            for seat_name in seat_map.keys()
                            if (raw := ticket_list.get(seat_name, {}).get("num")) is not None
                            and str(raw).strip() not in invalid_vals
                        }
                        
                        if available_seats:
                            start_time, end_time = None, None
                            start_keys = ['出发时间', '始发时间', '发车时间', '出发时刻']
                            end_keys = ['到达时间', '终到时间', '到站时间', '到达时刻']
                            for k in start_keys:
                                if v := train_info.get(k): start_time = str(v).strip(); break
                            for k in end_keys:
                                if v := train_info.get(k): end_time = str(v).strip(); break

                            ticket_info = {
                                "日期": query_date, "车次": actual_train_no, "出发地": dep_city,
                                "到达地": arr_city, "出发时间": start_time or "", "到达时间": end_time or "",
                                "余票席位": available_seats
                            }
                            print(f"  --> ✅ [查询成功] {dep_city} -> {arr_city} | 车次 {actual_train_no} 发现余票")
                            return ticket_info
                        else:
                            print(f"  --> 🟡 [已找到但无票] {dep_city} -> {arr_city} | 车次 {actual_train_no} 当前无余票。")
                            return None
                return None
            except Exception as e:
                print(f"  --> 🚨 [查询异常] {dep_city} -> {arr_city}: {e}")
            return None
        # --- 任务结束时，令牌会自动归还 ---

    all_results = []
    try:
        async with streamablehttp_client(url_with_auth) as (read, write, _):
            async with ClientSession(read, write) as session:
                print("正在初始化与 MCP 服务器的连接...")
                await session.initialize()
                print("✅ 服务器初始化成功！\n")
                
                tasks = []
                for dep_city in departure_cities:
                    for arr_city in arrival_cities:
                        if dep_city == arr_city: continue
                        tasks.append(query_single_pair(session, dep_city, arr_city))
                
                print(f"准备并发执行 {len(tasks)} 个查询任务...")
                results_from_gather = await asyncio.gather(*tasks, return_exceptions=True)
                
                for res in results_from_gather:
                    if isinstance(res, Exception):
                        print(f"  --> 🚨 [任务执行异常]: {res}")
                    elif res:
                        all_results.append(res)
    except ExceptionGroup as eg:
        print(f"❌ 连接或初始化 MCP 会话时发生并发错误。包含 {len(eg.exceptions)} 个子异常:")
        for i, e in enumerate(eg.exceptions):
            print(f"  - 子异常 {i+1}: {type(e).__name__} - {e}")
    except Exception as e:
        print(f"❌ 连接或初始化 MCP 会话失败: {type(e).__name__} - {e}")
        return []

    return all_results

# --- 使用示例 (与之前相同) ---
async def main():
    my_api_key = "cff0ac8f-51ef-4b46-ae98-3b1875b59ddb"
    target_date = "2025-10-01"
    target_train = ["G2691","G2694"]

    # ["潮汕", "普宁", "陆丰", "惠东南", "深圳北", "东莞南", "博罗北", "河源北", "龙川西", "和平北", "龙南东", "赣州西", "新干东", "南昌西", "庐山", "黄梅东", "太湖南", "安庆", "铜陵", "无为", "合肥南"]
    departure_points = ["潮汕", "普宁", "陆丰", "惠东南", "深圳北"]
    arrival_points = [  "太湖南", "安庆","郑州东"]

    if not my_api_key or "YOUR_API_KEY" in my_api_key:
        print("❌ 错误：请在 main 函数中设置您的 my_api_key。")
        return

    target_trains = [t.strip() for t in (target_train if isinstance(target_train, (list, tuple, set)) else [target_train]) if t]

    results = await find_available_tickets_for_train(
        my_api_key, target_trains, departure_points, arrival_points, target_date
    )

    print("\n---------- ✅ 查询任务完成，结果汇总如下 ----------")
    if results:
        sorted_results = sorted(results, key=lambda x: (x['出发地'], x['车次']))
        print(json.dumps(sorted_results, ensure_ascii=False, indent=4))
    else:
        print("所有查询组合均未发现符合条件的余票。")

if __name__ == '__main__':
    asyncio.run(main())